"""
This file contains the class definition for the ProcessedMeasurement class.
This class is used to store the processed measurement data from the car.
created by Tao Xu, 2024.5.1
"""

import numpy as np
import time
from bins.tool import *
import math
import numpy.lib.recfunctions as rfn
import logging

logger = logging.getLogger(__name__)


class ProcessedMeasurement:
    def __init__(self):
        self.dtype = np.dtype([
            ('x', np.float32),
            ('y', np.float32),
            ('h', np.float32),
            ('v', np.float32),
            ('t', np.float32)
        ])
        self.nearest_index = 0
        self.current_state = None
        self.match_state = None
        self.predicted_step = 40
        self.time_step = 0.05
        self.traj_point = None
        self.x_ref = None
        self.y_ref = None
        self.yaw_ref = None
        self.v_ref = None
        self.need_predict_traj = None
        self.calculate_current_state = None
        self.time = 0

    def init_traj(self, traj) -> bool:
        """
        初始化车辆轨迹
        通过传入的轨迹串初始化当前轨迹串
        """
        if self.traj_point is not None:
            self.traj_point = None

        if traj is None:
            logger.debug("No init trajectory data.")
            return False
        if len(traj) == 0:
            logger.debug("No init trajectory data.")
            return False
        dtype = np.dtype([
            ('x', np.float32),
            ('y', np.float32),
            ('h', np.float32),
            ('v', np.float32)
        ])
        if type(traj) is list:
            if len(traj[0]) == 4:
                self.traj_point = np.zeros(len(traj), dtype=self.dtype)
                traj = np.array(traj, dtype=dtype)
            elif len(traj[0]) == 5:
                self.traj_point = np.array(traj, dtype=self.dtype)
                return True
            else:
                logger.debug("Error in init_traj for len != 4 or 5.")
                return False
        if type(traj) is np.ndarray:
            if len(traj[0]) == 4:
                default_time = 0.0
                self.traj_point = np.zeros(traj.shape, dtype=self.dtype)
                for i in dtype.names:
                    self.traj_point[i] = traj[i]
                self.traj_point['t'] = default_time
                return True
            elif len(traj[0]) == 5:
                self.traj_point = traj
                return True
            else:
                logger.debug("Error in init_traj for len != 4 or 5.")
                return False
        logger.debug("Error in init_traj for type error.")
        return False

    def is_empty(self) -> bool:
        """
        判断是否为空
        """
        if self.traj_point is None or len(self.traj_point) == 0:
            return True
        return False


    def update_traj(self, traj) -> bool:
        """
        更新车辆轨迹
        根据传入轨迹串更新当前轨迹串
        通过匹配传入轨迹点的初始位置与当前轨迹点的最近位置，更新当前轨迹串
        """
        if traj is None:
            logger.debug(f"No update trajectory data.")
            return False
        if len(traj) == 0:
            logger.debug(f"No update trajectory data.")
            return False
        if self.traj_point is None:
            self.init_traj(traj)
            return True
        if type(traj) is list:
            traj = np.array(traj, dtype=self.dtype)
        nearest_idx = self.find_nearest_point(traj[0])
        self.traj_point = self.np.concatenate((self.traj_point[:nearest_idx], traj))
        return True

    def update_state(self, current_state) -> bool:
        """
        更新车辆状态
        根据传入的车辆当前状态计算车辆轨迹最近点
        输出匹配的MPC轨迹点
        """
        if current_state is None:
            logger.debug(f"No update state data.")
            return False
        if len(current_state) == 0:
            logger.debug(f"No update state data.")
            return False
        if type(current_state) is list:
            # 需要特殊的异常处理来判断长度
            if len(current_state) == 4:
                current_state.append(0)
            current_ = np.zeros(1, dtype=self.dtype)
            current_[0] = tuple(current_state)
            current_state = current_

        if len(current_state) != 5:
            current_ = np.zeros(1, dtype=self.dtype)
            for i in self.dtype.names:
                if i == 't':
                    current_[i] = 0
                    continue
                current_[i] = current_state[i]
            current_state = current_

        # 更新当前车辆状态
        try:
            self.current_state = current_state
            if self.processing_traj(current_state):
                return True
        except:
            logger.debug(f"Error in update_state for processing traj.")
        return False

    def processing_traj(self, current_state) -> bool:
        """
        处理车辆轨迹
        通过当前匹配的最近点,计算参考点
        """
        if self.traj_point is None:
            return False
        if len(self.traj_point) == 0:
            return False
        nearest_idx = self.find_nearest_point(current_state, self.nearest_index)
        if nearest_idx == -1:
            return False
        self.nearest_index = nearest_idx
        return True

    def find_nearest_point(self, current_state, begin_index: int = None) -> int:
        """
        查找最近点
        通过当前车辆状态，查找当前轨迹串中最近的点
        """
        if self.traj_point is None:
            return -1
        if len(self.traj_point) == 0:
            return -1
        if begin_index is None:
            begin_index = 0

        traj_point = self.traj_point[begin_index:]

        distance = np.sqrt((traj_point['x'] - current_state['x']) ** 2 +
                           (traj_point['y'] - current_state['y']) ** 2)
        nearest_idx = np.argmin(distance) + begin_index

        if distance[nearest_idx - begin_index] < 1e-2:  # 是一个小阈值，表示重合
            if nearest_idx + 1 < len(self.traj_point):  # 确保有下一个点
                return nearest_idx + 1  # 返回下一个点
            else:
                return nearest_idx  # 如果已经是最后一个点，则返回该点

        return nearest_idx

    def calculate_ref_point(self) -> bool:
        """
        计算参考轨迹点
        通过当前匹配的最近点，计算参考轨迹点，更新需要预测的轨迹点
        """
        if self.traj_point is None:
            return False
        if len(self.traj_point) == 0:
            return False
        if self.current_state is None:
            return False
        if len(self.current_state) == 0:
            return False
        # nearest_idx = self.find_nearest_point(self.current_state, self.nearest_index)
        nearest_idx = self.find_nearest_point(self.current_state, 0)
        if nearest_idx == -1:
            return False
        match_point, _ = project_point_to_line_segment((float(self.current_state['x'][0]), float(self.current_state['y'][0])),
                                                       self.traj_point[nearest_idx - 1], self.traj_point[nearest_idx])
        self.match_state = np.array(match_point, dtype=self.dtype)
        self.nearest_index = nearest_idx
        if self.transformed_map_traj(self.traj_point[nearest_idx:]):
            return True
        return False

    def transformed_map_traj(self, traj) -> bool:
        """
        转换地图轨迹
        将大地坐标转化为当前车身坐标系
        """
        if traj is None:
            logger.debug(f"No transformed trajectory data for no traj.")
            return False
        if len(traj) == 0:
            logger.debug(f"No transformed trajectory data for traj is 0.")
            return False
        if self.current_state is None:
            logger.debug(f"No transformed trajectory data for no current state.")
            return False
        if len(self.current_state) == 0:
            logger.debug(f"No transformed trajectory data for current state is 0.")
            return False
        traj_state = traj.copy()
        if len(traj_state) > self.predicted_step:
            traj_state = traj_state[:self.predicted_step]
        else:
            append_state = np.array([traj_state[-1]] * (self.predicted_step - len(traj_state)), dtype=self.dtype)
            traj_state = np.concatenate((traj_state, append_state))

        # traj_state[0] = self.match_state

        traj_state['x'] -= self.match_state['x']
        traj_state['y'] -= self.match_state['y']
        new_x = traj_state['x'] * np.cos(self.match_state['h']) - traj_state['y'] * np.sin(self.match_state['h'])
        new_y = traj_state['x'] * np.sin(self.match_state['h']) + traj_state['y'] * np.cos(self.match_state['h'])
        traj_state['x'] = new_x
        traj_state['y'] = new_y
        for i in range(len(traj_state)):
            delta_h = float(traj_state[i]['h']) - float(self.match_state['h'])
            traj_state[i]['h'] = delta_h

            # if self.match_state['h'] > 340 and traj_state[i]['h'] < 20:
            #     traj_state[i]['h'] += 360
            # elif traj_state[i]['h'] > 340 and self.match_state['h'] < 20:
            #     traj_state[i]['h'] -= 360
            # traj_state[i]['h'] -= self.match_state['h']
        # traj_state['h'] -= self.match_state['h']

        self.need_predict_traj = traj_state

        self.calculate_current_state = self.current_state
        self.calculate_current_state['x'] -= self.match_state['x']
        self.calculate_current_state['y'] -= self.match_state['y']
        delta_h = self.calculate_current_state['h'] - self.match_state['h']
        self.calculate_current_state['h'] = delta_h
        new_x_ = self.calculate_current_state['x'] * np.cos(self.match_state['h']) - self.calculate_current_state['y'] * np.sin(self.match_state['h'])
        new_y_ = self.calculate_current_state['x'] * np.sin(self.match_state['h']) + self.calculate_current_state['y'] * np.cos(self.match_state['h'])
        self.calculate_current_state['x'] = new_x_
        self.calculate_current_state['y'] = new_y_
        s = []
        for i in range(len(traj_state)):
            if i == 0:
                s.append(float(np.sqrt((new_x_ - traj_state['x'][i]) ** 2 + (new_y_ - traj_state['y'][i]) ** 2)))
                continue
            s.append(float(np.sqrt((traj_state['x'][i] - traj_state['x'][i - 1]) ** 2 +
                     (traj_state['y'][i] - traj_state['y'][i - 1]) ** 2) + s[i - 1]))
            # s.append(float(np.sqrt((traj_state['x'][i] - traj_state['x'][i - 1]) ** 2 +
            #                        (traj_state['y'][i] - traj_state['y'][i - 1]) ** 2)))
        self.need_predict_traj = rfn.append_fields(self.need_predict_traj, 's', s, usemask=False)
        return True

    def get_predict_traj(self):
        """
        获取预测轨迹
        """
        return self.need_predict_traj

